﻿using System;
using System.Collections.Generic;
using System.Diagnostics.Contracts;

namespace O1.Data.Traversal
{
    [ContractVerification(true)]
    public class PartitionedCellSet
    {
        private const int InitialBuckets = 13;

        private readonly int initialCapacity;

        private Entry[] entries;

        private int[] buckets;

        private int count;

        public PartitionedCellSet(int initialCapacity)
        {
            Contract.Ensures(16 <= this.initialCapacity);
            Contract.Ensures(this.entries != null);
            Contract.Ensures(this.entries.Length == InitialBuckets);
            Contract.Ensures(this.buckets != null);
            Contract.Ensures(this.buckets.Length == InitialBuckets);
            
            // TODO: How to associate mutiple depths with the same
            // cell to traverse? deduplicate....and associate with a depth.
            // can sorted list work here? or too expensive?
            // try naive first, then match.
            this.entries = new Entry[InitialBuckets];
            var pBuckets = new int[InitialBuckets];
            for (int i = 0; i < pBuckets.Length; ++i)
            {
                pBuckets[i] = -1;
            }

            this.buckets = pBuckets;
            this.initialCapacity = initialCapacity < 16 ? 16 : initialCapacity;
        }

        [CLSCompliant(false)]
        public ushort MaxTrunkId { get; private set; }

        [CLSCompliant(false)]
        public List<uint> this[ushort trunkId]
        {
            get
            {
                var pBuckets = this.buckets;
                var bucket = trunkId % pBuckets.Length;
                
                Contract.Assert(0 <= bucket);

                var pEntries = this.entries;
                for (int i = pBuckets[bucket]; i >= 0; i = pEntries[i].Next)
                {
                    Contract.Assume(i < pEntries.Length);

                    if (pEntries[i].TrunkId == trunkId)
                    {
                        return pEntries[i].CellIds;
                    }
                }

                return null;
            }
        }

        [CLSCompliant(false)]
        public void Add(uint cellId, ushort trunkId)
        {
            Contract.Ensures(this.buckets != null);
            Contract.Ensures(this.buckets.Length > 0);

            var pBuckets = this.buckets;
            var bucket = trunkId % pBuckets.Length;
            
            Contract.Assert(0 <= bucket);

            var pEntries = this.entries;
            for (int i = pBuckets[bucket]; i >= 0; i = pEntries[i].Next)
            {
                Contract.Assume(i < pEntries.Length);

                var entry = pEntries[i];
                if (entry.TrunkId == trunkId)
                {
                    Contract.Assume(entry.CellIds != null);
                    entry.CellIds.Add(cellId);
                    return;
                }
            }

            var index = this.count;
            if (index == this.entries.Length)
            {
                var newSize = (this.count * 2).RoundToPositivePrime();
                Contract.Assert(this.count < newSize);
                
                pBuckets = new int[newSize];
                for (int i = 0; i < pBuckets.Length; i++)
                {
                    pBuckets[i] = -1;
                }

                pEntries = new Entry[newSize];
                Array.Copy(this.entries, 0, pEntries, 0, this.count);

                for (int i = 0; i < this.count; i++)
                {
                    var entry = pEntries[i];
                    bucket = entry.TrunkId % newSize;
                    entry.Next = pBuckets[bucket];
                    pBuckets[bucket] = i;
                }

                this.buckets = pBuckets;
                this.entries = pEntries;
                bucket = trunkId % newSize;
            }

            // Use the existing list in the event the set is cleared and re-used.
            var items = pEntries[index].CellIds;
            if (items == null)
            {
                items = new List<uint>(this.initialCapacity);
                pEntries[index].CellIds = items;
            }

            items.Add(cellId);

            this.count++;
            pEntries[index].TrunkId = trunkId;
            pEntries[index].Next = pBuckets[bucket];
            pBuckets[bucket] = index;

            this.MaxTrunkId = Math.Max(this.MaxTrunkId, trunkId);
        }

        public void Clear()
        {
            if (this.buckets == null)
            {
                return;
            }

            for (int i = 0; i < this.count; i++)
            {
                if (this.entries[i].CellIds != null)
                {
                    this.entries[i].CellIds.Clear();
                }

                this.buckets[i] = -1;
            }

            for (int i = this.count; i < this.buckets.Length; i++)
            {
                this.buckets[i] = -1;
            }

            this.count = 0;
            this.MaxTrunkId = 0;
        }

        [ContractInvariantMethod]
        [System.Diagnostics.CodeAnalysis.SuppressMessage("Microsoft.Performance", "CA1822:MarkMembersAsStatic", Justification = "Required for code contracts.")]
        private void ObjectInvariant()
        {
            Contract.Invariant(this.entries != null);
            Contract.Invariant(this.entries.Length > 0);
            Contract.Invariant(this.buckets != null);
            Contract.Invariant(this.buckets.Length > 0); 
            Contract.Invariant(this.count >= 0);
            Contract.Invariant(this.count < this.entries.Length);
            Contract.Invariant(this.initialCapacity > 0);
            Contract.Invariant(this.buckets.Length == this.entries.Length);
        }
        
        private struct Entry
        {
            public ushort TrunkId;

            public int Next;

            public List<uint> CellIds;
        }
    }
}
